/**
    Copyright (c) 2010 yakiimo02
    Distributed under the New BSD License.
    See included license.txt or http://www.yakiimo3d.com/NewBSDLicense.txt
**/

#include "DXUT.h"
#include "MeasureLuminanceCS.h"
#include "DXHelper.h"

#include <vector>

// Constant buffer layout for transferring data to the CS
struct CB_CS
{
    UINT    param[4];
};

template <class T>
void SWAP( T* &x, T* &y )
{
    T* temp = x;
    x = y;
    y = temp;
}

/**
*/
MeasureLuminanceCS::MeasureLuminanceCS() : m_pReduceTo1DCS( NULL ), m_pReduceToSingleCS( NULL ), m_pBufferReduction0( NULL ), m_pBufferReduction1( NULL ), 
                         m_pReductionRV0( NULL ), m_pReductionRV1( NULL ), m_pReductionUAView0( NULL ), m_pReductionUAView1( NULL ), m_pcbCS( NULL ),
                         m_nReduceTo1DX_( 0 ), m_nReduceTo1DY_( 0 ), m_pGPUtoCPUBuf( NULL ), m_fLogAvgLum( 0.0f ), m_fLumMax( 0.0f )
{
}

/**
*/
MeasureLuminanceCS::~MeasureLuminanceCS()
{
}

/**
*/
HRESULT MeasureLuminanceCS::OnCreateDevice( ID3D11Device* pd3dDevice )
{
    HRESULT hr = S_OK;

    // compile compute shaders
    V_RETURN( CompileShaders( pd3dDevice ) );

    // Setup constant buffers
    D3D11_BUFFER_DESC Desc;
    Desc.Usage = D3D11_USAGE_DYNAMIC;
    Desc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
    Desc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
    Desc.MiscFlags = 0;    
    Desc.ByteWidth = sizeof( CB_CS );
    V_RETURN( pd3dDevice->CreateBuffer( &Desc, NULL, &m_pcbCS ) );

    return hr;
}

/**
*/
HRESULT MeasureLuminanceCS::OnResizedSwapChain( ID3D11Device* pd3dDevice, const int nWidth, const int nHeight )
{
    HRESULT hr = S_OK;

    OnReleasingSwapChain();

    // ReduceTo1DCS has threadgroup size of 8x8 threads
    m_nReduceTo1DX_ = int(ceil(nWidth / 8.0f));
    m_nReduceTo1DY_ = int(ceil(nHeight / 8.0f));

    // Create two buffers for ping-ponging in the reduction operation used for calculating luminance
    D3D11_BUFFER_DESC DescBuffer;
    ZeroMemory( &DescBuffer, sizeof(D3D11_BUFFER_DESC) );
    DescBuffer.BindFlags = D3D11_BIND_UNORDERED_ACCESS | D3D11_BIND_SHADER_RESOURCE;
    DescBuffer.ByteWidth = m_nReduceTo1DX_ * m_nReduceTo1DY_ * (sizeof(float)*2);
    DescBuffer.MiscFlags = D3D11_RESOURCE_MISC_BUFFER_STRUCTURED;
    DescBuffer.StructureByteStride = sizeof(float)*2;
    DescBuffer.Usage = D3D11_USAGE_DEFAULT;
    V_RETURN( pd3dDevice->CreateBuffer( &DescBuffer, NULL, &m_pBufferReduction0 ) );
    V_RETURN( pd3dDevice->CreateBuffer( &DescBuffer, NULL, &m_pBufferReduction1 ) );

    // create uavs 
    V_RETURN( pd3dDevice->CreateUnorderedAccessView( m_pBufferReduction0, NULL, &m_pReductionUAView0 ) );    
    V_RETURN( pd3dDevice->CreateUnorderedAccessView( m_pBufferReduction1, NULL, &m_pReductionUAView1 ) );    

    // create shader resource views
    V_RETURN( CreateBufferSRV( pd3dDevice, m_pBufferReduction0, &m_pReductionRV0 ) );
    V_RETURN( CreateBufferSRV( pd3dDevice, m_pBufferReduction1, &m_pReductionRV1 ) );

    // buffer to transferring gpu calculated log luminance average and luminance max to the cpu
    D3D11_BUFFER_DESC desc = DescBuffer;
    desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
    desc.Usage = D3D11_USAGE_STAGING;
    desc.BindFlags = 0;
    desc.MiscFlags = 0;
    V_RETURN( pd3dDevice->CreateBuffer( &desc, NULL, &m_pGPUtoCPUBuf ) );

    return hr;    
}

/**
*/
HRESULT MeasureLuminanceCS::CompileShaders( ID3D11Device* pd3dDevice )
{
    HRESULT hr = S_OK;
    ID3DBlob* pBlob = NULL;

    // Create the shaders
    V_RETURN( CompileShaderFromFile( L"ReduceTo1DCS.hlsl", "CSMain", "cs_5_0", &pBlob ) );   
    V_RETURN( pd3dDevice->CreateComputeShader( pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pReduceTo1DCS ) );
    SAFE_RELEASE( pBlob );

    V_RETURN( CompileShaderFromFile( L"ReduceToSingleCS.hlsl", "CSMain", "cs_5_0", &pBlob ) );
    V_RETURN( pd3dDevice->CreateComputeShader( pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pReduceToSingleCS ) );
    SAFE_RELEASE( pBlob );

    return hr;
}

/**
*/
void MeasureLuminanceCS::OnReleasingSwapChain()
{
    SAFE_RELEASE( m_pBufferReduction0 );
    SAFE_RELEASE( m_pBufferReduction1 );
    SAFE_RELEASE( m_pReductionRV0 );
    SAFE_RELEASE( m_pReductionRV1 );
    SAFE_RELEASE( m_pReductionUAView0 );
    SAFE_RELEASE( m_pReductionUAView1 );
    SAFE_RELEASE( m_pGPUtoCPUBuf );
}

/**
*/
void MeasureLuminanceCS::OnDestroyDevice()
{
    OnReleasingSwapChain();
    SAFE_RELEASE( m_pReduceTo1DCS );
    SAFE_RELEASE( m_pReduceToSingleCS );
    SAFE_RELEASE( m_pcbCS );
}

/**
    Measure the average luminance of the rendered skybox in CS path
*/
HRESULT MeasureLuminanceCS::DoMeasureLuminance( ID3D11DeviceContext* pd3dImmediateContext, int nWidth, int nHeight, ID3D11ShaderResourceView* pInputSRV )
{    
    HRESULT hr = S_OK;

    // First CS pass, reduce the render target texture into a 1D buffer
    {
        ID3D11ShaderResourceView* aRViews[ 1 ] = { pInputSRV };
        CB_CS cbCS = { m_nReduceTo1DX_, m_nReduceTo1DY_, nWidth, nHeight };
        RunComputeShader( pd3dImmediateContext, 
                          m_pReduceTo1DCS, 
                          1, aRViews,
                          m_pcbCS, &cbCS, sizeof(cbCS),
                          m_pReductionUAView0,
                          m_nReduceTo1DX_, m_nReduceTo1DY_, 1 ); 

        // Turn on this and set a breakpoint right after the call to Map to see what's written to g_pBufferReduction0
#if 0
        ID3D11Buffer* pDebugBuf = CreateAndCopyToDebugBuf( DXUTGetD3D11Device(), pd3dImmediateContext, g_pBufferReduction0 );
        D3D11_MAPPED_SUBRESOURCE MappedResource;     
        V( pd3dImmediateContext->Map( pDebugBuf, 0, D3D11_MAP_READ, 0, &MappedResource ) );
        pd3dImmediateContext->Unmap( pDebugBuf, 0 );
        SAFE_RELEASE( pDebugBuf );        
#endif
    }

    // Reduction CS passes, the reduction result will be in the first element of g_pTex1DReduction1
    {
        int dim = m_nReduceTo1DX_*m_nReduceTo1DY_;
        int nNumToReduce = dim;
        dim = int( ceil(dim/128.0f) );
        if ( nNumToReduce > 1 )
        {
            for (;;)
            {
                ID3D11ShaderResourceView* aRViews[ 1 ] = { m_pReductionRV0 };
                CB_CS cbCS = { nNumToReduce, dim, 0, 0 };
                RunComputeShader( pd3dImmediateContext,
                    m_pReduceToSingleCS,
                    1, aRViews,
                    m_pcbCS, &cbCS, sizeof(cbCS),
                    m_pReductionUAView1, 
                    dim, 1, 1 );                               

                nNumToReduce = dim;
                dim = int( ceil(dim/128.0f) );

                if ( nNumToReduce == 1 )
                    break;

                SWAP( m_pBufferReduction0, m_pBufferReduction1 );
                SWAP( m_pReductionUAView0, m_pReductionUAView1 );
                SWAP( m_pReductionRV0, m_pReductionRV1 );
            }
        } else
        {
            SWAP( m_pBufferReduction0, m_pBufferReduction1 );
            SWAP( m_pReductionUAView0, m_pReductionUAView1 );
            SWAP( m_pReductionRV0, m_pReductionRV1 );
        }
    }    

     pd3dImmediateContext->CopyResource( m_pGPUtoCPUBuf, m_pBufferReduction1 );
     D3D11_MAPPED_SUBRESOURCE MappedResource;     
     V( pd3dImmediateContext->Map( m_pGPUtoCPUBuf, 0, D3D11_MAP_READ, 0, &MappedResource ) );
     float* data  = reinterpret_cast< float* >( MappedResource.pData );
     m_fLogAvgLum = exp( ( data[ 0 ] / (nWidth*nHeight) ) );
     m_fLumMax = data[ 1 ];
     pd3dImmediateContext->Unmap( m_pGPUtoCPUBuf, 0 );

    // Turn on this and set a breakpoint right after the call to Map to see what is in g_pBufferReduction1 and g_pBufferReduction0
#if 0
    ID3D11Buffer* pDebugBuf = CreateAndCopyToDebugBuf( DXUTGetD3D11Device(), pd3dImmediateContext, g_pBufferReduction1 );
    D3D11_MAPPED_SUBRESOURCE MappedResource;     
    V( pd3dImmediateContext->Map( pDebugBuf, 0, D3D11_MAP_READ, 0, &MappedResource ) );
    pd3dImmediateContext->Unmap( pDebugBuf, 0 );
    SAFE_RELEASE( pDebugBuf );

    pDebugBuf = CreateAndCopyToDebugBuf( DXUTGetD3D11Device(), pd3dImmediateContext, g_pBufferReduction0 );    
    V( pd3dImmediateContext->Map( pDebugBuf, 0, D3D11_MAP_READ, 0, &MappedResource ) );
    pd3dImmediateContext->Unmap( pDebugBuf, 0 );
    SAFE_RELEASE( pDebugBuf );
#endif

    return hr;
}

/**    
    Helper function which makes CS invocation more convenient
*/
void MeasureLuminanceCS::RunComputeShader( ID3D11DeviceContext* pd3dImmediateContext,
                       ID3D11ComputeShader* pComputeShader,
                       UINT nNumViews, ID3D11ShaderResourceView** pShaderResourceViews, 
                       ID3D11Buffer* pCBCS, void* pCSData, DWORD dwNumDataBytes,
                       ID3D11UnorderedAccessView* pUnorderedAccessView,
                       UINT X, UINT Y, UINT Z )
{
    HRESULT hr = S_OK;
    
    pd3dImmediateContext->CSSetShader( pComputeShader, NULL, 0 );
    pd3dImmediateContext->CSSetShaderResources( 0, nNumViews, pShaderResourceViews );
    pd3dImmediateContext->CSSetUnorderedAccessViews( 0, 1, &pUnorderedAccessView, NULL );
    if ( pCBCS )
    {
        D3D11_MAPPED_SUBRESOURCE MappedResource;
        V( pd3dImmediateContext->Map( pCBCS, 0, D3D11_MAP_WRITE_DISCARD, 0, &MappedResource ) );
        memcpy( MappedResource.pData, pCSData, dwNumDataBytes );
        pd3dImmediateContext->Unmap( pCBCS, 0 );
        ID3D11Buffer* ppCB[1] = { pCBCS };
        pd3dImmediateContext->CSSetConstantBuffers( 0, 1, ppCB );
    }

    pd3dImmediateContext->Dispatch( X, Y, Z );

    ID3D11UnorderedAccessView* ppUAViewNULL[1] = { NULL };
    pd3dImmediateContext->CSSetUnorderedAccessViews( 0, 1, ppUAViewNULL, NULL );

    ID3D11ShaderResourceView* ppSRVNULL[3] = { NULL, NULL, NULL };
    pd3dImmediateContext->CSSetShaderResources( 0, 3, ppSRVNULL );
    ID3D11Buffer* ppBufferNULL[1] = { NULL };
    pd3dImmediateContext->CSSetConstantBuffers( 0, 1, ppBufferNULL );
}
